#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Thu Mar 21 20:43:36 2019
generate wanted alps lattice
this time hofstadter model ladder lattice
@author: bb
"""
class Vertex(object):
    # a vertex object that contains vertex ID, coordinates
    def __init__(self,ID):
        assert type(ID)==int,'ID has to be an integer.'
        self.ID = ID
        self.type = 0
        if ID%2==0: # if even
            self.coordinate=((ID-2)/2,1)
        else: # if odd
            self.coordinate=((ID-1)/2,0)
    def get_string(self):
        return '<VERTEX id="'+str(self.ID)+'" type="'+str(self.type)+'"><COORDINATE>'\
    +str(self.coordinate[0])+' '+str(self.coordinate[1])+'</COORDINATE></VERTEX>'

class Edge(object):
    Edge_ID = 1
    # an edge object that contains edge ID, source-target vertex ID, coordinate and type
    def __init__(self,vertexID1,vertexID2):
        assert type(vertexID1)==int and type(vertexID2)==int,'ID has to be an integer.'
        assert vertexID1<vertexID2,'vertexID1 should be smaller than vertexID2'
        
        self.ID = Edge.Edge_ID
        self.source = vertexID1
        self.target = vertexID2
        source_target_legnumber = self.get_source_target_legnumber()
        if source_target_legnumber[0]==source_target_legnumber[1]:
            self.type = source_target_legnumber[0]
            self.vector = '0 1'
        else:
            self.type = 0
            self.vector = '1 0'        
        Edge.Edge_ID += 1
    def get_source_target_legnumber(self):
        legnumbers = []
        for vertexID in [self.source,self.target]:
            if vertexID%2==0: # if even
                legnumbers.append(vertexID/2)
            else: # if odd
                legnumbers.append((vertexID+1)/2)
        return legnumbers
    def get_string(self):
        return '<EDGE source="'+str(self.source)+'" target="'+str(self.target)+'" id="'\
    +str(self.ID)+'" type="'+str(self.type)+'" vector="'+self.vector+'"/>'

def generate_hofstadter_model_ladder_lattice(L):
    assert type(L)==int,'L must be integer.'
#    L = 10; ## ladder的leg数，所以格点数应该翻倍
    IDs = range(1,2*L+1)
    vertexes = [Vertex(ID) for ID in IDs] # generate vertexes
    
    edges1 = []  # for [1 3 5 7...]
    list1 = [2*i-1 for i in range(1,L+1)] 
    for i in range(len(list1)-1):
        vertexID1 = list1[i]
        vertexID2 = list1[i+1]
        edges1.append(Edge(vertexID1,vertexID2))
    
    edges2 = [] # for [2 4 6 8...]
    list2 = [2*i for i in range(1,L+1)] 
    for i in range(len(list2)-1):
        vertexID1 = list2[i]
        vertexID2 = list2[i+1]
        edges2.append(Edge(vertexID1,vertexID2))
    
    list1 = [2*i-1 for i in range(1,L+1)] 
    edges3 = [Edge(i,i+1) for i in list1] # for [1 2],[3 4]...
    
    edges = edges1+edges2+edges3 #generate edges
    
    
    ## generate lattice xml file
    f=open('hofstadter_model_ladder_lattice_L%s.xml'%(L),'w')
    f.write('<LATTICES>\n')
    f.write('<GRAPH name="hofstadter_model_ladder_lattice" dimension="2" vertices="%s" edges="%s">\n'%(len(vertexes),len(edges)))
    for vertex in vertexes:
        f.write(vertex.get_string()+'\n')
    for edge in edges:
        f.write(edge.get_string()+'\n')
    f.write('</GRAPH>\n')
    f.write('</LATTICES>\n')
    f.close()
    
    lattice_library = 'hofstadter_model_ladder_lattice_L%s.xml'%(L)
    lattice_name = 'hofstadter_model_ladder_lattice'
    return (lattice_library,lattice_name)




















